package DecisionTree;

import com.google.common.base.Predicate;
import com.google.common.collect.*;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;

import java.io.File;

import static org.apache.commons.math3.util.FastMath.log;

/**
 * 作者: LDL
 * 功能说明: 决策树算法(ID3,C4.5)
 * 创建日期: 2015/6/3 17:40
 */
public class DecisionTree {

    public static void main(String[] args) throws IOException {
        //System.out.println(chooseBestFeatureToSplit(getDataSetFromFile("lenses.txt")));
        //System.out.println(spliteDataSet(getDataSetFromFile("lenses.txt"), 0, "pre"));
        ArrayList<String> labels = Lists.newArrayList("age", "prescript", "astigmatic", "tearRate");
        createTree(getDataSetFromFile("lenses.txt"),labels);


    }

    /**
     * 计算香农熵
     *@return
     */
    public static float calcShannonEnt(ArrayList<ArrayList<String>> dataSet){

        int totalSize = dataSet.size();
        float shannonEnt = 0;
        Multiset<String> labelSet = HashMultiset.create();
        for (int i = 0; i < dataSet.size(); i++) {
            ArrayList<String> data = dataSet.get(i);
            labelSet.add(data.get(data.size() - 1));
        }
        for(String key : labelSet.elementSet()){
            float prob = (float)labelSet.count(key)/totalSize;
            shannonEnt -= prob * log(prob,2);
        }
        return shannonEnt;
    }

    /**
     * 从文集读取
     * @param fileName
     * @return
     * @throws IOException
     */
    public static ArrayList<ArrayList<String>> getDataSetFromFile(String fileName) throws IOException {
        String path = DecisionTree.class.getClassLoader().getResource(fileName).getPath();
        BufferedReader reader = new BufferedReader(new FileReader(new File(path)));
        String line = null;
        ArrayList<ArrayList<String>> dataSet = Lists.newArrayList();
        while((line=reader.readLine())!=null){
            ArrayList<String> lineData = Lists.newArrayList();
            for(String str : line.split("\t")){
                lineData.add(str);
            }
            dataSet.add(lineData);
        }
        return dataSet;
    }

    /**
     * 按照给定的特征值来划分数据集
     * @param dataSet 带划分数据集
     * @param axis 给定的特征
     * @param value 特征的返回值
     * @return
     */
    public static ArrayList<ArrayList<String>> spliteDataSet(ArrayList<ArrayList<String>> dataSet, final int axis, final Object value){

        /**
         * 过滤条件
         */
        Predicate<ArrayList<String>> predicate = new Predicate<ArrayList<String>>() {
            @Override
            public boolean apply(ArrayList<String> input) {
                if (input.get(axis).equals(value)) {
                    return true;
                }
                return false;
            }
        };
        ArrayList<ArrayList<String>> retData = Lists.newArrayList(FluentIterable.from(dataSet).filter(predicate));
        for(ArrayList<String> data : dataSet){
            data.remove(value);
        }
        return retData;
    }

    /**
     * 根据信息增益选择最优划分属性的下标
     * @param dataSet
     * @return
     */
    public static int chooseBestFeatureToSplit(ArrayList<ArrayList<String>> dataSet){
        int index = -1;
        int size = dataSet.get(0).size()-1;
        int count = dataSet.size();
        float baseEntropy = calcShannonEnt(dataSet);
        float bestInfoGain = 0;
        for (int i = 0; i < size; i++) {
            ArrayList<ArrayList<String>> copyData = new ArrayList<>();
            for (ArrayList<String> data : dataSet){
                ArrayList<String> d = new ArrayList<>(data);
                copyData.add(d);
            }
            HashSet<String> set = Sets.newHashSet();
            float newEntropy = 0;
            for (ArrayList<String> data : copyData){
                set.add(data.get(i));
            }

            for(String key : set){
                ArrayList<ArrayList<String>> subDataSet = spliteDataSet(copyData,i,key);
                float prob = (float)subDataSet.size()/count;
                newEntropy += prob * calcShannonEnt(subDataSet);
            }
            if(bestInfoGain>newEntropy){
                bestInfoGain = newEntropy;
                index = i;
            }
        }
        return index;
    }


    public static void createTree(ArrayList<ArrayList<String>> dataSet,ArrayList<String> labels){

        if(dataSet.get(0).size()==1){
            System.out.println("子节点："+dataSet.get(0).get(dataSet.get(0).size()-1));
            return;
        }
        HashSet<String> tset = Sets.newHashSet();
        for (int i = 0; i < dataSet.size(); i++) {
            ArrayList<String> data = dataSet.get(i);
            tset.add(data.get(data.size()-1));
        }
        if(tset.size()==1){
            System.out.println("子节点："+tset.toString());
            return;
        }
        int index = chooseBestFeatureToSplit(dataSet);
        if(index==-1){
            System.out.println("子节点："+dataSet.get(0).get(dataSet.get(0).size()-1));
            return;
        }
        System.out.println("当前根节点标签是：" + labels.get(index));
        HashSet<String> set = Sets.newHashSet();
        for (int i = 0; i < dataSet.size(); i++) {
            ArrayList<String> data = dataSet.get(i);
            set.add(data.get(index));
        }
        labels.remove(index);
        for(String key : set){
            System.out.println("特征："+key);
            ArrayList<ArrayList<String>> subDataSet = spliteDataSet(dataSet,index,key);
            createTree(subDataSet,labels);
        }
    }

    /*public static String majorityCnt(ArrayList<ArrayList<String>> dataSet){
    }*/

}
